# %%
from torchvision import transforms
from PIL import Image
import os
import torch
import torchvision
import os
from data_loader import show_images
from train import find_best_gpu
find_best_gpu()
root_dir='/home/jinyue/work/pcy241003/data/edeme/'
save_dir='/home/jinyue/work/pcy241003/checkpoints'
transform = transforms.Compose([
  transforms.Grayscale(num_output_channels=1),  # 转换为单色图片（1通道）
    transforms.Resize((512,
    512)),                # 调整大小为 512x512
    transforms.ToTensor()                         # 转换为 PyTorch Tensor
])


def show1(train_set):
    # 示例展示
    n = 5
    imgs = []
    for i in range(n):
        f, l = train_set[i
]

        # 将 feature 和 label 转换为 NumPy 数组
        img_feature = f.numpy().transpose(1,
2,
0)
        img_label = l.numpy().transpose(1,
2,
0)

        imgs.append(img_feature)
        imgs.append(img_label)

    # 假设你有一个显示图像的函数
    show_images(imgs, n,
2)
def save_checkpoint(model, optimizer, epoch, save_path):
    checkpoint = {
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'epoch': epoch
}
    torch.save(checkpoint, save_path)
    print(f"Checkpoint saved at {save_path}")

# %%
import os
import torch
import torchvision
from torch.utils.data import random_split, TensorDataset

def read_images(root_dir):
    features = []
    labels = []

    # 读取文件并限制读取数量
    label_dir = os.path.join(root_dir, 'label')
    pic_dir = os.path.join(root_dir, 'raw')

    for i, fname in enumerate(os.listdir(label_dir)):

        lab_path = os.path.join(label_dir, fname)
        pic_path = os.path.join(pic_dir, fname)

        if os.path.isfile(lab_path) and os.path.isfile(pic_path):
            # 读取图片并转换为 FloatTensor
            feature_image = torchvision.io.read_image(pic_path).float() / 255.0
            label_image = torchvision.io.read_image(lab_path).float() / 255.0

            features.append(feature_image)
            labels.append(label_image)

    return features, labels

# 读取图像数据
features, labels = read_images(root_dir)

# 转换为 TensorDataset
features_tensor = torch.stack(features)
labels_tensor = torch.stack(labels)
dataset = TensorDataset(features_tensor, labels_tensor)

# 数据集划分
total_size = len(dataset)
train_ratio, val_ratio, test_ratio = 0.6,0.2,0.2
train_size = int(train_ratio * total_size)
val_size = int(val_ratio * total_size)
test_size = total_size - train_size - val_size

# 随机划分训练集、验证集和测试集
train_set, val_set, test_set = random_split(dataset,
[train_size, val_size, test_size
])

show1(train_set)


# %%
from UNetResNet50 import UNetResNet50
# 定义UNet模型，使用ResNet-50作为骨干网络
import torch.nn as nn
import torch.optim as optim


# 超参数设置
epochs = 100
learning_rate = 1e-5
batch_size = 2  # 根据显存情况调整
num_classes = 1  # 假设是二分类问题
device = torch.device("cuda" )

# 初始化模型、损失函数和优化器
# model = UNetResNet50(num_classes=num_classes).to(device)
import segmentation_models_pytorch as smp

model = smp.Unet(
    encoder_name="resnet50",        # choose encoder, e.g. mobilenet_v2 or efficientnet-b7
    encoder_weights="imagenet",     # use `imagenet` pre-trained weights for encoder initialization
    in_channels=3,                  # model input channels (1 for gray-scale images,
    classes=1,                      # model output channels (number of classes in your dataset)
).to(device)
criterion = nn.BCEWithLogitsLoss()  # 二分类问题
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 训练过程
for epoch in range(epochs):
    model.train()
    running_loss = 0.0
    for images, labels in train_set:
        images, labels = images.to(device), labels.to(device)

        # 检查输入维度并确保是 4D，如果需要，可以使用 unsqueeze(0) 来扩展 batch 维度
        if images.dim() == 3:  # 如果输入是 3D (channels, height, width)
            images = images.unsqueeze(0)  # 将 batch_size 维度添加到 0 维度
        if labels.dim() == 3:
            labels = labels.unsqueeze(0)

        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)

        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
    print(f"Epoch [{epoch+1}/{epochs}], Loss: {running_loss/len(train_set)}")
    save_path =save_dir+ f"/model_epoch_{epoch}.pth"

    save_checkpoint(model, optimizer, epoch, save_path)
'''
def pltre():
    plt.figure(figsize=(8,
6))
    plt.subplot(2,
2,
1)
    plt.plot(range(1, epoch + 2), train_accuracies, label='Train Accuracy')
    plt.plot(range(1, epoch + 2), val_accuracies, label='Val Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.title('Accuracy over Epochs')
    plt.legend()
    plt.savefig(f'/media/disk/01drive/07qiulin/pj1/graphics/_{epoch + 1
}_Accuracy.png')
    plt.close()
    '''

# %%

# 验证模型
model.eval()
with torch.no_grad():
    val_loss = 0.0
    for images, labels in val_set:
        images, labels = images.to(device), labels.to(device)
                # 检查输入维度并确保是 4D，如果需要，可以使用 unsqueeze(0) 来扩展 batch 维度
        if images.dim() == 3:  # 如果输入是 3D (channels, height, width)
            images = images.unsqueeze(0)  # 将 batch_size 维度添加到 0 维度
        if labels.dim() == 3:
            labels = labels.unsqueeze(0)
        outputs = model(images)
        loss = criterion(outputs, labels)
        val_loss += loss.item()

    print(f"Validation Loss: {val_loss/len(val_set)}")


# %%
import torch
def calculate_iou(y_true, y_pred):
    # Convert to binary
    y_true = y_true.int()
    y_pred = y_pred.int()

    intersection = torch.sum((y_true & y_pred).float())
    union = torch.sum((y_true | y_pred).float())

    iou_score = intersection / union
    return iou_score.item()

def calculate_dice(y_true, y_pred):
    # Convert to binary
    y_true = y_true.int()
    y_pred = y_pred.int()

    intersection = torch.sum((y_true & y_pred).float())
    dice_score = (2 * intersection) / (torch.sum(y_true.float()) + torch.sum(y_pred.float()))
    return dice_score.item()

def evaluate_segmentation(y_true, y_pred):
    # Flatten tensors
    y_true_flat = y_true.view(-1)
    y_pred_flat = y_pred.view(-1)

    # Calculate IoU
    iou = calculate_iou(y_true_flat, y_pred_flat)

    # Calculate Dice coefficient
    dice = calculate_dice(y_true_flat, y_pred_flat)

    # Calculate Precision
    precision = torch.sum((y_true_flat & y_pred_flat).float()) / torch.sum(y_pred_flat.float())

    # Calculate Recall
    recall = torch.sum((y_true_flat & y_pred_flat).float()) / torch.sum(y_true_flat.float())

    # Calculate F1-score
    f1 = 2 * (precision * recall) / (precision + recall + 1e-8)

    # Calculate Pixel Accuracy
    accuracy = torch.sum((y_true_flat == y_pred_flat).float()) / len(y_true_flat)

    return {
        "IoU": iou,
        "Dice Coefficient": dice,
        "Precision": precision.item(),
        "Recall": recall.item(),
        "F1-score": f1.item(),
        "Pixel Accuracy": accuracy.item()
    }

# 示例：生成假数据进行评价
# 假设 y_true 和 y_pred 是二进制的分割图（0 表示背景，1 表示前景）
y_true = torch.tensor([[0, 1, 1], [0, 1, 0], [0, 0, 0]])
y_pred = torch.tensor([[0, 1, 0], [0, 1, 1], [0, 0, 0]])

# 计算评价指标
metrics = evaluate_segmentation(y_true, y_pred)
print(metrics)
# %%

# 测试集上的评估
def evaluate_model(model, loader):
    model.eval()
    with torch.no_grad():
        metrics = {
    "PA": 0,
    "CPA": 0,
    "MPA": 0,
    "IoU": 0,
    "MIoU": 0
    }
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
                        # 检查输入维度并确保是 4D，如果需要，可以使用 unsqueeze(0) 来扩展 batch 维度
            if images.dim() == 3:  # 如果输入是 3D (channels, height, width)
                images = images.unsqueeze(0)  # 将 batch_size 维度添加到 0 维度
            if labels.dim() == 3:
                labels = labels.unsqueeze(0)
            outputs = model(images)
            # 计算像素准确率 (PA)
            pred = outputs.argmax(dim=1)
            correct = (pred == labels).float()
            pa = correct.mean().item()
            metrics[
    "PA"
    ] += pa

            # 计算类别像素准确率 (CPA) 和平均像素准确率 (MPA)
            num_classes = outputs.size(1)
            cpa = []
            for cls in range(num_classes):
                cls_pred = pred == cls
                cls_true = labels == cls
                if cls_true.sum() > 0:
                    cls_pa = (cls_pred & cls_true).float().sum() / cls_true.sum()
                    cpa.append(cls_pa.item())
            metrics[    "CPA"        ] += sum(cpa)
            metrics[    "MPA"] += sum(cpa) / len(cpa) if cpa else 0

            # 计算交并比 (IoU) 和平均交并比 (MIoU)
            iou = []
            for cls in range(num_classes):
                intersection = ((pred == cls) & (labels == cls)).float().sum()
                union = ((pred == cls) | (labels == cls)).float().sum()
                if union > 0:
                    cls_iou = intersection / union
                    iou.append(cls_iou.item())
            metrics[    "IoU"] += sum(iou)
            metrics[    "MIoU"] += sum(iou) / len(iou) if iou else 0
        return metrics



# 将模型移动到可用的设备上
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 加载 state_dict（这是一个 OrderedDict）
state_dict = torch.load('unet_resnet50_model.pth')

# 将 state_dict 加载到模型中
model.load_state_dict(state_dict)


test_metrics = evaluate_model(model, test_set)
print(f"Test Metrics: {test_metrics}")



# %%



